# -- coding: utf-8 --
'''
预处理训练/测试集
'''
from __future__ import division
from __future__ import print_function
import glob
import define
import os
import sys
from openslide import open_slide, OpenSlideUnsupportedFormatError
import multiresolutionimageinterface as mir
from PIL import Image, ImageDraw
from itertools import product
from datetime import datetime
import Prostate_preprocessing.file_ops as file_ops
import numpy as np
import cv2
import time
import pdb
import re
import tensorflow as tf
import random
import argparse

class WSI(object):

	def __init__(self, DATA_SET_FLAG):
		'''
		ROI提取中:
		close_kernel：控制白色区域的扩增，值越大，白色越能连成片
		open_kernel：控制黑色块，用以消除零散的白色片区
		lower_red：控制深色上限
		upper_red：控制浅色下限
		'''
		self._DATA_SET_FLAG = DATA_SET_FLAG
		self._patch_size = 256         #所要提取的样本尺寸
		self._workingLevel = 1         #样本提取所在的分辨层

		self._overview_level = 4       #缩略图的分辨层
		self._scale_factor = pow(2, self._overview_level-self._workingLevel)  #层间缩放因子

		self._negative_patch_index = 0
		self._p_negative_patch_index = 0
		self._positive_patch_index = 0
		self._positive_patch_found = 0
		self._lower_red = np.array([20, 50, 20]) #对浅色穿刺样本建议改为：np.array([10, 10, 10])
		self._upper_red = np.array([200, 150, 200])
		self._open_kernel = np.ones((30, 30), dtype=np.uint8)
		self._close_kernel = np.ones((25, 25), dtype=np.uint8)
		self._TissueMask = None #Image类型
		self._patch_distribution_map = None #Array类型	

	def _int64_feature(self, value):
		'''
		被saveAs_TFRecords()调用，解析int64型数据
		'''
		return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))

	def _bytes_feature(self, value):
		'''
		被saveAs_TFRecords()调用，解析bytes型数据
		'''
		return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))

	def GetPatch(self, RAWImg, start_w, start_h, workingLevel, windowShape):
		'''
		从WSI对象中提取patch

		Args：
		start_w, start_h：左上角起始点坐标
		workingLevel：执行样本提取所用的level
		windowShape：提取区域大小
		'''
		patch = np.array(RAWImg.read_region((start_w, start_h), workingLevel, windowShape)) 
		return patch;

	def read_wsi(self, wsi_path, mask_path=None):
		'''
		把原图读为level-4的RGB
		对于tumor，原图和对应mask一起读
		'''
		try:
			self._wsi_filename = file_ops.get_filename_from_path(wsi_path)
			self._wsi_image = open_slide(wsi_path)
			
			wsi_ov_size = self._wsi_image.level_dimensions[self._overview_level]

			rgba_image = self.GetPatch(self._wsi_image, 0, 0, self._overview_level, wsi_ov_size)
			self._rgb_image = cv2.cvtColor(rgba_image, cv2.COLOR_BGRA2BGR)

			if mask_path:
				self._mask_image = open_slide(mask_path)
				rgba_mask = self.GetPatch(self._mask_image, 0, 0, self._overview_level, wsi_ov_size)         
				self._rgb_mask = cv2.cvtColor(rgba_mask, cv2.COLOR_BGRA2BGR)*255   #由(0,1)转为(0,255)
				#self._mask_image、self._rgb_mask为array类型

		except OpenSlideUnsupportedFormatError:
			print('Exception: OpenSlideUnsupportedFormatError')
			return False

		return True

	def get_image_contours(self, cont_img, obj_image, line_color='green', line_thick = 5):
		'''
		画出组织轮廓，给出各块区域的bbox

		args：
		cont_img：基于cont_img计算轮廓及bbox
		obj_image：在obj_image上绘制轮廓及bbox等
		line_color：bbox或轮廓颜色，取值为(red/green/blue/white)
		line_thick：bbox或轮廓线宽，若需要填充区域，取line_thick=-1
		'''
		color_map = {'red':(0,0,255), 'green':(0,255,0), 'blue':(255,0,0), 'white':(255,255,255)}
		assert line_color in color_map, 'line_color not in color_map, you should give a value in (red/green/blue/white)'
		line_color = color_map[line_color]
		_, contours, _ = cv2.findContours(cont_img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)  
		bounding_boxes = []
		for c in contours:
			x,y,w,h = cv2.boundingRect(c)    
			bounding_box = [x,y,x+w,y+h]
			bounding_boxes.append(bounding_box)
		img_with_contour = cv2.drawContours(obj_image, contours, -1, line_color, line_thick)    #mask的轮廓采用填充方式
		return obj_image, bounding_boxes

	def extract_patches_mask(self, DATA_SET_FLAG, bounding_boxes, mask_path):
		if DATA_SET_FLAG=='Train':
			label_txt_file_path = define.PROSTATE_PATCHES_TRAIN_TXT +self._wsi_filename + '.txt'
			patches_dir = define.PROSTATE_PATCHES_TRAIN_TUMOR_POS
			patch_distribution_dir = define.TRAIN_PATCH_DISTRIBUTION_DIR
		elif DATA_SET_FLAG=='Test':
			label_txt_file_path = define.PROSTATE_PATCHES_VALIDATION_TXT +self._wsi_filename + '.txt'
			patches_dir = define.PROSTATE_PATCHES_VALIDATION_POS
			patch_distribution_dir = define.VALIDATION_PATCH_DISTRIBUTION_DIR
		label_txt_file = open(label_txt_file_path ,'w')
		mask_pixel_rate = 0.7
		tissue_pixel_rate = 0.4
		win_size = self._patch_size * pow(2, self._workingLevel) 	#计算指定的self._patch_size在level-0上的大小
		mag_factor = pow(2, self._overview_level) 					#计算从level-0到self._overview_level（这里为4）的缩放因子
		step = win_size 											#不重叠取样，win_size = step
		lower_red = np.array([150, 10, 50])							#基于patch进行hsv转换
		upper_red = np.array([200, 200, 240])		
		total_found = 0
		filename = file_ops.get_filename_from_path(mask_path)
		print ('Start extracting Tumor patches from %s' % filename)
		for i, bounding_box in enumerate(bounding_boxes):			#逐个bbox进行扫描提取
			print ('No.%d of Tumor ROIs to extract... %s\n' % (i+1,filename),end='\r')
			b_x_start = int(bounding_box[0]) * mag_factor			#bbox在self._workingLevel获取，需要转换回level-0
			b_y_start = int(bounding_box[1]) * mag_factor
			b_x_end = int(bounding_box[2]) * mag_factor
			b_y_end = int(bounding_box[3]) * mag_factor
			if b_x_end-b_x_start < win_size:						#一系列越界判断处理
				X = [(b_x_end-b_x_start)/2]
			else:
				X = np.arange(b_x_start + win_size, b_x_end - win_size, step)  #patch中心点坐标列表
			if b_y_end-b_y_start < win_size:
				Y = [(b_y_end-b_y_start)/2]
			else:
				Y = np.arange(b_y_start + win_size, b_y_end - win_size, step)
			total_patches = len(X) * len(Y)							#根据中心点列表估计待检测的patch总数
			print ('--%d patches to detect--' % total_patches)
			m_iter = t_iter = 0 									#临时变量，统计处理进度
			interval = total_patches/20

			_gray_mask = Image.fromarray(self._rgb_mask).convert('L')
			patch_level_area = pow(self._patch_size/pow(2,(self._overview_level-self._workingLevel)), 2)		#patch在self._overview_level上的面积
			patch_win_area = pow(self._patch_size, 2)															#patch取样面积,按self._patch_size的平方算

			_rgb_image_tmp = self._rgb_image 	#_rgb_image_tmp临时变量，用于生成取样分布图
			for x, y in list(product(X, Y)):	#逐个patch判断
				if t_iter > interval:			#进度统计
					t_iter = 0
					print ('%.2f %% finished' %(m_iter*100/total_patches))
				x_corner = x - win_size/2 		#patch在self._overview_level上的左上角左边
				y_corner = y - win_size/2
				patch_size = (self._patch_size, self._patch_size)			

				crop_size = (int(x_corner/mag_factor), int(y_corner/mag_factor), int((x_corner+win_size)/mag_factor), int((y_corner+win_size)/mag_factor))  #在self._overview_level上的裁剪面积
				mask = _gray_mask.crop(crop_size) 						#mask在self._overview_level同一位置上的裁剪块
				TargetMask_array = np.array(mask)
				white_pixel_cnt_gt = cv2.countNonZero(TargetMask_array) #计算这一patch对应的mask上同一位置的白色像素（标注区）占比

				if white_pixel_cnt_gt > (patch_level_area * mask_pixel_rate): 		#0.70：当patch中标注区像素占比达到70%
					patch = self.GetPatch(self._wsi_image, int(x_corner), int(y_corner), self._workingLevel, patch_size)	#在self._workingLevel（不一定为level-0）上截取，但其他参数均为level-0上的值，注意！
					patch_array = np.array(patch)
					patch_hsv = cv2.cvtColor(patch_array, cv2.COLOR_BGR2HSV)
					mask_patch = cv2.inRange(patch_hsv, lower_red, upper_red)
					white_pixel_cnt = cv2.countNonZero(mask_patch)

					if white_pixel_cnt > (patch_win_area * tissue_pixel_rate):	#二次判断，为了剔除标注区中存在的组织空洞区域（即排除标注区中的空洞，由于标注区轮廓较为粗糙，无法排除其中的组织空洞）
						self._patch_distribution_map = cv2.rectangle(_rgb_image_tmp, crop_size[:2], crop_size[2:4], (0,0,255),5)	#记录取样位置（画带颜色的与win_size同等大小的矩形框）,array类型
						self._positive_patch_index += 1
						patch_save_path = patches_dir + define.PATCH_TUMOR_PREFIX + str(self._positive_patch_index) + '.jpg'
						cv2.imwrite(patch_save_path, patch)						#保存符合要求的patch
						patch_name = file_ops.get_filename_from_path(patch_save_path)
						 #记录条目到txt file，格式：WSI名称，patch命名，标签，x坐标，y坐标，patch存储路径
						label_txt_file.write(self._wsi_filename + "," + patch_name + "," + "1," + str(int(x)) + "," + str(int(y)) + "," + patch_save_path + "\n")
						self._positive_patch_found += 1
						total_found += 1
				m_iter += 1
				t_iter += 1
		label_txt_file.close()
		#注意：若需要提取非标注区，以下保存patch_distribution_map将被更新
		if total_found:	#对于一些及其浅色的WSI，有可能一个样本都提取不到
			patch_distribution_path = patch_distribution_dir + filename + "_distribution_map.jpg"
			self.save(patch_distribution_path, self._patch_distribution_map)
		print ('Finish extracting Tumor patches from %s, %d targets found' % (filename,total_found))
		return total_found
   # bounding_boxes 在overview level4 下取得的bounding boxes 
   # 将start_x, end_x, start_y ,end_y 转换成提取patch的 workingLevel 1 下的坐标b_x_start, b_x_end, b_y_start,b_y_end, 循环截取patch 
   # 做两次过滤， 1. 将patch的面积区域坐标转换对应到overview level4上的 tissueMask上的面积区域，如果非0区域占大于mask_pixel_rate =95% 那么继续 
   # 2. 将patch做GBR->HSV 转换，取patch_hsv mask图，计算非0面积，如果大于tissue_pixel_rate =25% 那么继续存储patch.
	def extract_patches_normal(self, DATA_SET_FLAG, bounding_boxes, wsi_path, equal_get_nums=None, Tnormal=False):
		if DATA_SET_FLAG=='Train':
			label_txt_file_path = define.PROSTATE_PATCHES_TRAIN_TXT +self._wsi_filename + '.txt'
			patches_dir =  define.PROSTATE_PATCHES_TRAIN_TUMOR_NEG if Tnormal else define.PROSTATE_PATCHES_TRAIN_NORMAL_NEG
			patch_distribution_dir = define.TRAIN_PATCH_DISTRIBUTION_DIR
		elif DATA_SET_FLAG=='Test':
			label_txt_file_path = define.PROSTATE_PATCHES_VALIDATION_TXT +self._wsi_filename + '.txt'
			patches_dir = define.PROSTATE_PATCHES_VALIDATION_NEG
			patch_distribution_dir = define.VALIDATION_PATCH_DISTRIBUTION_DIR	
		label_txt_file = open(label_txt_file_path,'a')
		mask_pixel_rate = 0.95 if Tnormal else 0.9
		tissue_pixel_rate = 0.25 if Tnormal else 0.1
		win_size = self._patch_size * pow(2, self._workingLevel)
		mag_factor = pow(2, self._overview_level)  #self._scale_factor
		step = win_size
		lower_red = np.array([150, 10, 50])
		upper_red = np.array([200, 200, 240])
		total_found = 0
		filename = file_ops.get_filename_from_path(wsi_path)
		if Tnormal:
			print ('Start extracting Tumor-Normal patches from %s' % filename)
		else:
			print ('Start extracting Normal patches from %s' % filename)
		for i, bounding_box in enumerate(bounding_boxes):
			if Tnormal:
				print ('No.%d of Tumor-Normal ROIs to extract... %s' % (i+1,filename))
			else:
				print ('No.%d of Normal ROIs to extract... %s' % (i+1,filename))
			b_x_start = int(bounding_box[0]) * mag_factor
			b_y_start = int(bounding_box[1]) * mag_factor
			b_x_end = int(bounding_box[2]) * mag_factor
			b_y_end = int(bounding_box[3]) * mag_factor
			if b_x_end-b_x_start < win_size:
				X = [(b_x_end-b_x_start)/2]
			else:
				X = np.arange(b_x_start + win_size, b_x_end - win_size, step)
			if b_y_end-b_y_start < win_size:
				Y = [(b_y_end-b_y_start)/2]
			else:
				Y = np.arange(b_y_start + win_size, b_y_end - win_size, step)
			
			random_extract_list = list(product(X, Y))	#Normal patch为随机抽取
			random.shuffle(random_extract_list)
			total_patches = equal_get_nums if equal_get_nums else len(random_extract_list)
			print ('--%d patches to detect--' % len(random_extract_list)) 	#detech所有patch，直到找到total_patches个patch为止
			m_iter = t_iter = 0
			interval = total_patches/20
			patch_level_area = pow(self._patch_size/pow(2,(self._overview_level-self._workingLevel)), 2)  #
			patch_win_area = pow(self._patch_size, 2)
			_rgb_image_tmp = self._rgb_image
			for x, y in random_extract_list:
				if t_iter > interval:
					t_iter = 0
					print ('%.2f %% finished' %(m_iter*100/total_patches))
				x_corner = x - win_size/2
				y_corner = y - win_size/2
				patch_size = (self._patch_size, self._patch_size)

				crop_size = (int(x_corner/mag_factor), int(y_corner/mag_factor), int((x_corner+win_size)/mag_factor), int((y_corner+win_size)/mag_factor))
				mask = self._TissueMask.crop(crop_size)
				TissueMask_array = np.array(mask)
				white_pixel_cnt_gt = cv2.countNonZero(TissueMask_array)

				if white_pixel_cnt_gt > (patch_level_area * mask_pixel_rate):
					patch = self.GetPatch(self._wsi_image, int(x_corner), int(y_corner), self._workingLevel, patch_size)                   
					patch_array = np.array(patch)
					patch_hsv = cv2.cvtColor(patch_array, cv2.COLOR_BGR2HSV)
					mask_patch = cv2.inRange(patch_hsv, lower_red, upper_red)
					white_pixel_cnt = cv2.countNonZero(mask_patch)
					if white_pixel_cnt > (patch_win_area * tissue_pixel_rate):
						self._patch_distribution_map = cv2.rectangle(_rgb_image_tmp, crop_size[:2], crop_size[2:4], (0,255,0),5)
						if Tnormal:
							self._p_negative_patch_index += 1
						else:
							self._negative_patch_index += 1
						if Tnormal:
							patch_save_path = patches_dir + define.PATCH_NORMAL_PREFIX + str(self._p_negative_patch_index) + '.jpg'
						else:
							patch_save_path = patches_dir + define.PATCH_NORMAL_PREFIX + str(self._negative_patch_index) + '.jpg'
						cv2.imwrite(patch_save_path, patch)
						patch_name = file_ops.get_filename_from_path(patch_save_path)
						label_txt_file.write(self._wsi_filename + "," + patch_name + "," + "0," + str(int(x)) + "," + str(int(y)) + "," + patch_save_path + "," + "\n")
						total_found += 1
				m_iter += 1
				t_iter += 1
				if equal_get_nums and (total_found == total_patches): 	#找到足够的patch就返回
					break
			if equal_get_nums and (total_found == total_patches):
					break
		if total_found:
			patch_distribution_path = patch_distribution_dir + filename + "_distribution_map.jpg"
			self.save(patch_distribution_path, self._patch_distribution_map)
			print ("Patch distribution map has been saved to: " + patch_distribution_path)
		label_txt_file.close()
		if Tnormal:
			print ('Finish extracting Tumor-Normal patches from %s, %d targets found.' % (filename,total_found))
		else:
			print ('Finish extracting Normal patches from %s, %d targets found.' % (filename,total_found))
		return total_found

	def find_roi_n_extract_patches(self, DATA_SET_FLAG, wsi_path, mask_path=None, equal_get_nums=None):
		hsv_image = cv2.cvtColor(self._rgb_image, cv2.COLOR_BGR2HSV)
		mask = cv2.inRange(hsv_image, self._lower_red, self._upper_red)
		image_close = Image.fromarray(cv2.morphologyEx(np.array(mask), cv2.MORPH_CLOSE, self._close_kernel))
		image_open = Image.fromarray(cv2.morphologyEx(np.array(image_close), cv2.MORPH_OPEN, self._open_kernel))
		contour_wsi, tissue_bounding_boxes = self.get_image_contours(np.array(image_open), self._rgb_image) #基于image_open计算轮廓并在原WSI上显示
		gray_mask = 0 #默认为不存在Tumor mask，若存在，该变量将被更新
		if mask_path:
			gray_mask = cv2.cvtColor(self._rgb_mask, cv2.COLOR_BGR2GRAY)
			_, mask_bounding_boxes = self.get_image_contours(gray_mask, self._rgb_mask, line_color='white', line_thick = -1) #基于gray_mask计算标注区的bbox
		TissueMask = np.array(image_open) - gray_mask
		self._TissueMask = Image.fromarray(TissueMask)
		_, tissue_bounding_boxes = self.get_image_contours(TissueMask, self._rgb_image) #基于TissueMask计算正常区的bbox，这里综合考虑了全Normal及Tumor WSI的Normal区域
		filename = file_ops.get_filename_from_path(wsi_path)
		if DATA_SET_FLAG=='Train':
			contour_wsi_save_path = define.TRAIN_TISSUE_CONTOUR_DIR + filename + "_TissueMask.jpg"
		else:
			contour_wsi_save_path = define.VALIDATION_TISSUE_CONTOUR_DIR + filename + "_TissueMask.jpg"
		self.save(contour_wsi_save_path, contour_wsi)
		print ("TissueContour has been saved to: " + contour_wsi_save_path)
		if mask_path:
			positive_patch_found = self.extract_patches_mask(DATA_SET_FLAG, mask_bounding_boxes, mask_path)
			# self.extract_patches_normal(DATA_SET_FLAG, tissue_bounding_boxes, wsi_path, equal_get_nums=positive_patch_found, Tnormal=True) #抽取Tumor WSI中的非标注区域
			return positive_patch_found
		else:
			negative_patch_found = self.extract_patches_normal(DATA_SET_FLAG, tissue_bounding_boxes, wsi_path, equal_get_nums=equal_get_nums)
			return negative_patch_found

	def display(self, contour_wsi, factor=0.15):
		'''
		显示

		args：
		contour_rgb：图
		factor：缩放因子
		'''
		contour_wsi = cv2.resize(contour_wsi, (0, 0), fx=factor, fy=factor)
		cv2.imshow(self._wsi_filename, np.array(contour_wsi))
		cv2.waitKey(0)
			
	def draw_bbox(self, contour, bounding_boxes):
		'''
		绘制bbox

		args：
		contour：底图
		bounding_boxes：bbox列表
		'''
		for i, bounding_box in enumerate(bounding_boxes):
			top_left = tuple(bounding_box[:2])
			bot_right = tuple(bounding_box[2:4])
			contour_bbox = cv2.rectangle(contour,top_left, bot_right,(0,0,255),5)
		self.display(contour_bbox)

	def save(self, save_path, contour_rgb):
		'''
		保存图

		args：
		save_path：路径
		contour_rgb：图
		'''
		contour_rgb = cv2.resize(contour_rgb, (0, 0), fx=0.40, fy=0.40)
		cv2.imwrite(save_path, np.array(contour_rgb))

	def equal_shffule(self, DATA_SET_FLAG='Train'):
		'''
		用于打乱WSI提取的patch的txt文件的各行，也就是打乱patch的顺序，针对tumor WSI，能够让同一个txt file中的正负样本混合
		#####保留功能，本程序没有启用#####
		'''
		if DATA_SET_FLAG=='Train':
			tumor_txt_file_path = define.PROSTATE_PATCHES_TRAIN_TXT +self._wsi_filename + '_tumor' + '.txt'
			normal_txt_file_path = define.PROSTATE_PATCHES_TRAIN_TXT +self._wsi_filename + '_normal' + '.txt'
			shuffle_txt_file_path = define.PROSTATE_PATCHES_TRAIN_TXT +self._wsi_filename + '_shuffle' + '.txt'
		elif DATA_SET_FLAG=='Test':
			tumor_txt_file_path = define.PROSTATE_PATCHES_VALIDATION_TXT +self._wsi_filename + '_tumor' + '.txt'
			normal_txt_file_path = define.PROSTATE_PATCHES_VALIDATION_TXT +self._wsi_filename + '_normal' + '.txt'
			shuffle_txt_file_path = define.PROSTATE_PATCHES_VALIDATION_TXT +self._wsi_filename + '_shuffle' + '.txt'
		tumor_txt_file = open(tumor_txt_file_path ,'r')
		normal_txt_file = open(normal_txt_file_path ,'r')
		shuffle_txt_file = open(shuffle_txt_file_path ,'w')

		Tumor_rec = tumor_txt_file.readlines()
		Normal_rec = normal_txt_file.readlines()
		shuffle_rec = []
		shuffle_rec.extend(Tumor_rec)
		shuffle_rec.extend(Normal_rec)
		random.shuffle(shuffle_rec)
		shuffle_txt_file.writelines(shuffle_rec)

		shuffle_txt_file.close()
		tumor_txt_file.close()
		normal_txt_file.close()

	def saveAs_TFRecords(self, DATA_SET_FLAG='Train'):
		'''
		用于读取WSI生成的txt file中的记录并生成TFRecords，程序中采取的处理方式是每个WSI对应一个TFRecords
		'''
		label_txt_file_name = self._wsi_filename + '.txt'
		if DATA_SET_FLAG=='Train':
			TFRecord_Dir = define.TRAIN_TFRECORDS_DIR
			txt_file_path = glob.glob(os.path.join(define.PROSTATE_PATCHES_TRAIN_TXT, label_txt_file_name))
		elif DATA_SET_FLAG=='Test':
			TFRecord_Dir = define.VALIDATION_TFRECORDS_DIR
			txt_file_path = glob.glob(os.path.join(define.PROSTATE_PATCHES_VALIDATION_TXT, label_txt_file_name))
		rec_cnt = 0
		shuffle_file = open(txt_file_path[0],'r')
		TFRcord_file = TFRecord_Dir + self._wsi_filename + '.tfrecords'     
		patchName =[]
		label = []
		Xcoor =[]
		Ycoor =[]
		path = []
		Raw_recs = shuffle_file.readlines()
		writer = tf.python_io.TFRecordWriter(TFRcord_file)
		for i in xrange(np.shape(Raw_recs)[0]):
			patchPath = Raw_recs[i].split(',')[5].split('\n')[0]
			image = Image.open(patchPath)
			image_raw = image.tobytes()
			label = int(Raw_recs[i].split(',')[2])
			Xcoor = int(Raw_recs[i].split(',')[3])
			Ycoor = int(Raw_recs[i].split(',')[4])
			Source = Raw_recs[i].split(',')[0]
			patchName=Raw_recs[i].split(',')[1]
			example = tf.train.Example(features=tf.train.Features(feature={
				'image_raw':self._bytes_feature(image_raw),
				'label':self._int64_feature(label),
				'Xcoor':self._int64_feature(Xcoor),
				'Ycoor':self._int64_feature(Ycoor),
				'Source':self._bytes_feature(Source),
				'patchName':self._bytes_feature(patchName),
				}))
			writer.write(example.SerializeToString())
			rec_cnt += 1
		print('file has been writen into TFRecords. %d records total'%rec_cnt) 
		writer.close()
		sys.stdout.flush()

def rename_data(DATA_SET_FLAG):
	'''
	重命名文件，一般为初始化的时候用，若是后续增加的文件需要注意命名编号是否能够对应，避免WSI,XML,MASK之间对应出错

	args：
	DATA_SET_FLAG：Train/Test ：对Train/Test数据集进行重命名
	'''
	if DATA_SET_FLAG=='Train':
		file_ops.rename_forDir(train_tumor_wsi_dir, 'Tumor', train_tumor_mask_dir, 
									define.TRAIN_TUMOR_XML_DIR, exclude_Dir)
	elif DATA_SET_FLAG=='Test':
		file_ops.rename_forDir(train_tumor_wsi_dir, 'Test', train_tumor_mask_dir, 
									define.VALIDATION_TUMOR_XML_DIR)

def run_on_tumor_data(DATA_SET_FLAG='Train', train_tumor_wsi_dir=define.TRAIN_TUMOR_WSI_DIR, train_tumor_mask_dir=define.TRAIN_TUMOR_MASK_DIR,wsi_type_patten=define.GLOB_WSI_TYPE_PATTEN, mask_type_patten=define.GLOB_MASK_TYPE_PATTEN):
	'''
	训练集预处理函数
	'''
	wsi = WSI(DATA_SET_FLAG)
	if DATA_SET_FLAG=='Train':
		wsi.wsi_paths = glob.glob(os.path.join(train_tumor_wsi_dir, wsi_type_patten))
		wsi.mask_paths = glob.glob(os.path.join(train_tumor_mask_dir, mask_type_patten))
	else:
		wsi.wsi_paths = glob.glob(os.path.join(test_tumor_wsi_dir, wsi_type_patten))
		wsi.mask_paths = glob.glob(os.path.join(test_tumor_mask_dir, mask_type_patten))
	wsi.wsi_paths.sort()
	wsi.mask_paths.sort()
	total_found = 0
	positive_patch_found = 0
	Total_WSI = len(wsi.wsi_paths)
	Total_MASK = len(wsi.mask_paths)
	#由于TestSet中mask和WSI数量不同，需要创建一个匹配字典进行对应
	maskdict = {}
	if Total_WSI == 0:
		print ('Failed to read Slides')
		return None
	for i in range(Total_MASK):
		mask_name = file_ops.get_filename_from_path(wsi.mask_paths[i])
		maskdict[mask_name] = wsi.mask_paths[i]
	for i in range(Total_WSI):
		wsi_path = wsi.wsi_paths[i]
		wsi_name = file_ops.get_filename_from_path(wsi_path)
		if maskdict.has_key(wsi_name):
			mask_path = maskdict[wsi_name]
		else:
			mask_path = None
			
		if os.path.exists(mask_path) and wsi.read_wsi(wsi_path, mask_path):
			print ('---Processing for No.%d WSI(Tumor), Total:%d, %.2f %%finished'%(i+1,Total_WSI,(i/Total_WSI)*100))
			positive_patch_found = wsi.find_roi_n_extract_patches(DATA_SET_FLAG, wsi_path, mask_path)
			
		elif wsi.read_wsi(wsi_path):
			print ('---Processing for No.%d WSI(Normal), Total:%d, %.2f %%finished'%(i+1,Total_WSI,(i/Total_WSI)*100))
			wsi.find_roi_n_extract_patches(DATA_SET_FLAG, wsi_path)  
		else:
			print ('Failed to read file, please check your file path:')
			print ('wsi_path:%s' %wsi_path)
			print ('mask_path:%s' %mask_path)
		# wsi.equal_shffule(DATA_SET_FLAG)
		wsi.saveAs_TFRecords(DATA_SET_FLAG)
		total_found += positive_patch_found
		print('Till now %d positive patches has been found\n'%total_found)
				
	else:
		print ('Processing for tumor_data has finished')
		print ("Extracting patches from WSI, total = %d, positive = %d, negetive = %d" \
					%(wsi._positive_patch_index + wsi._negative_patch_index, wsi._positive_patch_index, wsi._negative_patch_index))
		return total_found

def run_on_normal_data(DATA_SET_FLAG='Train', train_normal_wsi_dir=define.TRAIN_NORMAL_WSI_DIR, wsi_type_patten=define.GLOB_WSI_TYPE_PATTEN, positive_patch_found=None):
	'''
	训练集Normal WSI预处理函数

	args:
	positive_patch_found:对Tumor WSI预处理所提取到的Tumor patches数目，一般为run_on_tumor_data()的返回值。
						默认为None。使用这个参数，可以基于Tumor patches改变从Normal WSI中提取的patches总数。
	'''
	wsi = WSI(DATA_SET_FLAG)
	wsi.wsi_paths = glob.glob(os.path.join(train_normal_wsi_dir, wsi_type_patten))
	wsi.wsi_paths.sort()

	Total_WSI = len(wsi.wsi_paths)
	Total_big_WSI = 0
	#若指定了positive_patch_found，则按照平衡取样的方式，取样数（equal_get_nums）=Tumor patches总数/Nomal WSI的数量，其中
	#忽略小于300Mb的WSI（相当于不计穿刺样本），避免计算偏差太大
	for pth in wsi.wsi_paths:
		if (os.path.getsize(pth)/1024/1024)>300:
			Total_big_WSI+=1
	equal_get_nums = int(positive_patch_found/Total_big_WSI) if positive_patch_found else None

	for i in range(Total_WSI):
		wsi_path = wsi.wsi_paths[i]
		wsi_name = file_ops.get_filename_from_path(wsi_path)
		if wsi.read_wsi(wsi_path): #逐张读取
			print ('---Processing for No.%d normal WSI, Total:%d, %.2f %%finished'%(i+1,Total_WSI,(i/Total_WSI)*100))
			wsi.find_roi_n_extract_patches(DATA_SET_FLAG, wsi_path, equal_get_nums=equal_get_nums)  #提取
			wsi.saveAs_TFRecords(DATA_SET_FLAG)
		else:
			print ('Failed to read %s, please check your file path:'%wsi_name)
			print ('wsi_path:%s' %wsi_path)
	if Total_WSI == 0:
		print ('Failed to read Normal slides')
	else:
		print ('Processing for normal_data has finished')
		print ("Extracting patches from tumor WSI, total = %d, positive = %d, negetive = %d" \
					%(wsi._positive_patch_index + wsi._negative_patch_index, wsi._positive_patch_index, wsi._negative_patch_index))

def check_on_tumor_data(DATA_SET_FLAG='Train', train_tumor_wsi_dir=define.TRAIN_TUMOR_WSI_DIR, train_tumor_mask_dir=define.TRAIN_TUMOR_MASK_DIR, test_tumor_wsi_dir=define.VALIDATION_WSI_DIR, test_tumor_mask_dir=define.VALIDATION_TUMOR_MASK_DIR, wsi_type_patten=define.GLOB_WSI_TYPE_PATTEN, mask_type_patten=define.GLOB_MASK_TYPE_PATTEN):
	'''
	预处理过程需要对patch、TFRecords等进行编号，不可中断，否则需要重新运行。
	因此，预处理之前需要先对所有WSI及mask进行预加载试验，避免某些文件的损坏导致预处理过程被中断。
	执行此函数若遇到不能被正常读取的wsi或mask，自行找办法处理
	此函数作用于训练集Tumor及整个测试集，对于测试集的WSI会自动搜索是否存在对应的mask，存在则一并加载，否则只加载WSI

	args：
	DATA_SET_FLAG：Train/Test标记
	train_tumor_wsi_dir：训练集Tumor wsi存储位置
	train_tumor_mask_dir：训练集mask存储位置
	test_tumor_wsi_dir：测试集wsi存储位置
	test_tumor_mask_dir：测试集mask存储位置
	wsi_type_patten：wsi格式，此处为.ndpi
	mask_type_patten：mask格式，此处为.tif
	'''
	wsi = WSI(DATA_SET_FLAG)
	if DATA_SET_FLAG=='Train':
		wsi.wsi_paths = glob.glob(os.path.join(train_tumor_wsi_dir, wsi_type_patten))
		wsi.mask_paths = glob.glob(os.path.join(train_tumor_mask_dir, mask_type_patten))
	else:
		wsi.wsi_paths = glob.glob(os.path.join(test_tumor_wsi_dir, wsi_type_patten))
		wsi.mask_paths = glob.glob(os.path.join(test_tumor_mask_dir, mask_type_patten))
	wsi.wsi_paths.sort()
	wsi.mask_paths.sort()
	Total_WSI = len(wsi.wsi_paths)
	Total_MASK = len(wsi.mask_paths)
	#由于TestSet中mask和WSI数量不同，需要创建一个匹配字典进行对应
	maskdict = {}
	for i in range(Total_MASK):
		mask_name = file_ops.get_filename_from_path(wsi.mask_paths[i])
		maskdict[mask_name] = wsi.mask_paths[i]
	for i in range(Total_WSI):
		wsi_name = file_ops.get_filename_from_path(wsi.wsi_paths[i])
		if maskdict.has_key(wsi_name):
			mask_path = maskdict[wsi_name]
		else:
			mask_path = None
		if os.path.exists(mask_path):
			print ('[%d/%d]  checking on %s' %(i+1, Total_WSI, wsi_name))
			wsi.read_wsi(wsi.wsi_paths[i], mask_path)
		else:
			print ('[%d/%d]  checking on %s' %(i+1, Total_WSI, wsi_name))
			wsi.read_wsi(wsi.wsi_paths[i])
				
def check_on_normal_data(train_normal_wsi_dir=define.TRAIN_NORMAL_WSI_DIR, wsi_type_patten=define.GLOB_WSI_TYPE_PATTEN, positive_patch_found=None):
	'''
	预处理过程需要对patch、TFRecords等进行编号，不可中断，否则需要重新运行。
	因此，预处理之前需要先对所有WSI及mask进行预加载试验，避免某些文件的损坏导致预处理过程被中断。
	执行此函数若遇到不能被正常读取的wsi或mask，自行找办法处理
	此函数作用于训练集Normal

	args：
	train_normal_wsi_dir：训练集Normal wsi存储位置
	wsi_type_patten：wsi格式，此处为.ndpi
	mask_type_patten：mask格式，此处为.tif
	'''
	wsi = WSI('Train')
	wsi.wsi_paths = glob.glob(os.path.join(train_normal_wsi_dir, wsi_type_patten))
	wsi.wsi_paths.sort()
	Total_WSI = len(wsi.wsi_paths)
	for i in range(Total_WSI):
		wsi_path = wsi.wsi_paths[i]
		wsi_name = file_ops.get_filename_from_path(wsi_path)
		print ('[%d/%d]  checking on %s' %(i+1, Total_WSI, wsi_name))
		wsi.read_wsi(wsi_path) #逐张读取

def timecost(duration):
	'''
	时间格式转换：s → h:m:s
	
	Args：
	duration：秒数
	'''
	secs = duration%60
	hours = duration//3600
	minutes = duration//60-hours*60
	print ('Time cost: %dh %dmin %dsecs' % (hours,minutes,secs))

def xml2mask(WSI_dir=define.TRAIN_TUMOR_WSI_DIR, XML_dir=define.TRAIN_TUMOR_XML_DIR,output_mask_dir=define.TRAIN_TUMOR_MASK_DIR):
	'''
	将Tumor WSI的癌区标记XML文档转换为癌区识别掩膜，若癌区中央包含需要剔除的正常区域（最终输出为带空洞的掩膜），
	需要有优化版ASAP的加持。对于多分类情况对应的多色标记也可以用这一函数一次性处理
	需要特别注意的是：
	1、此函数需要配合ASAP源码一并修改，修改ASAP部分为：ASAP-1.7.3/annotation/AnnotationToMask.cpp → void AnnotationToMask::convert()
		修改后需要进行编译：
		步骤：$ cd release
			  $ cmake -D CMAKE_BUILD_TYPE=RELEASE -D CMAKE_INSTALL_PREFIX=/usr/opt/ ..  //这里后边有两个点不能省略
			  $ sudo make
			  $ sudo make install
	2、ASAP无法适应带group和不带group的annotation共存，修改ASAP源码后编译通过，但调用会出错（段错误）。
		所以标记的时候要么全部不属于group，要么全部属于，具体参照已有的XML文档。
	3、编译ASAP，出现OPENSLIDE_INCLUDE_DIR找不到，可直接打开Cmakecache.txt
		找到OPENSLIDE_INCLUDE_DIR,添加=/usr/include/openslide
	xml2mask需要较长时间处理，耐心等待

	Args：
	WSI_dir：Tumor WSI存储位置
	XML_dir：Tumor WSI对应的XML标记文档存储位置
	output_mask_dir：函数输出mask的存储位置
	'''
	reader = mir.MultiResolutionImageReader()
	wsi_paths = glob.glob(os.path.join(WSI_dir, define.GLOB_WSI_TYPE_PATTEN))
	wsi_paths.sort()
	xml_paths = glob.glob(os.path.join(XML_dir, define.GLOB_XML_TYPE_PATTEN))
	xml_paths.sort()
	Total_WSI = len(wsi_paths)
	Total_XML = len(xml_paths)
	assert Total_XML==Total_WSI, '标记文件数应于WSI文件数相等'
	for i in range(Total_XML):
		start_time = time.time()
		WSI_name = wsi_paths[i].split('/')[-1].split('.')[0]
		#加载Tumor WSI
		mr_image = reader.open(wsi_paths[i])
		annotation_list = mir.AnnotationList()
		xml_repository = mir.XmlRepository(annotation_list)
		xml_repository.setSource(xml_paths[i])
		xml_repository.load()
		annotation_mask = mir.AnnotationToMask()
		#根据在ASAP中的标记命名，最好一开始跟做标记的医生做好沟通，统一命名
		label_map = {'exclude': 0, 'Tumor': 1}
		output_path = output_mask_dir + '/' + WSI_name + '.tif'
		print ('[%d/%d]  converting mask of %s... ' %(i+1, Total_XML, WSI_name))
		annotation_mask.convert(annotation_list, output_path, mr_image.getDimensions(), mr_image.getSpacing(), label_map)
		timecost(time.time()-start_time)

def run_on_test_data():
	'''
	测试集预处理函数,使用run_on_tumor_data()处理，可以针对测试集正负样本并存的情况进行处理
	'''
	DATA_SET_FLAG = 'Test'
	run_on_tumor_data(DATA_SET_FLAG, define.VALIDATION_WSI_DIR, define.VALIDATION_TUMOR_MASK_DIR)

if __name__ == '__main__':
	parser = argparse.ArgumentParser(description='脚本参数说明')
	parser.add_argument('--dataset_type', type=str, default = None, help = "训练集 or 测试集？")
	parser.add_argument('--xml2mask', type=str, default = False, help = "是否执行生成mask？")
	parser.add_argument('--checkWSI', type=str, default = None, help = "是否检查WSI文件正确性？")
	parser.add_argument('--balance_extraction', type=str, default = True, help = "是否平衡抽取？")

	# ---------指定文件一级目录（其余子目录结构严格参照define.py）--------
	parser.add_argument('--train_wsi_dir', type=str, default = define.PUBLIC_TRAIN_DATASET, help = "训练集WSI目录")
	parser.add_argument('--test_wsi_dir', type=str, default = define.PUBLIC_VALIDATION_DATASET, help = "测试集WSI目录")
	args = parser.parse_args()

	assert args.dataset_type, '请指定数据集类型（Train/Test）'
	start_time = time.time()
	# ---------XML转mask--------
	# pdb.set_trace()
	if args.xml2mask=='True':
		xml2mask() 
	
	# ---------预加载 检查训练集WSI可读性--------
	if args.checkWSI=='True':
		if args.dataset_type == 'Train':
			print('Start checking on Train Tumor WSI, dir= %s.' %define.TRAIN_TUMOR_WSI_DIR)
			check_on_tumor_data()   
			print('Start checking on Train Normal WSI, dir= %s.' %define.TRAIN_NORMAL_WSI_DIR)  
			check_on_normal_data()
	# ---------预加载 检查测试集WSI可读性--------
		elif args.dataset_type == 'Test':
			print('Start checking on Test WSI, dir= %s.' %define.VALIDATION_WSI_DIR)
			check_on_tumor_data('Test', define.VALIDATION_WSI_DIR, define.VALIDATION_TUMOR_MASK_DIR)

	# ---------预处理训练集WSI--------
	if args.dataset_type == 'Train':
		if args.balance_extraction:
			positive_patch_found = run_on_tumor_data()  #预处理Tumor WSI并返回Tumor patches总数
			print('%d positive patches has been found.' %positive_patch_found)
			run_on_normal_data(positive_patch_found=positive_patch_found)  #预处理Normal WSI并根据Tumor patches总数平衡抽取
		else:
			run_on_tumor_data()
			run_on_normal_data()
	# ---------预处理测试集WSI--------
	if args.dataset_type == 'Test':
		run_on_test_data()
		print('Processing finish at %s'%time.strftime('%Y-%M-%d %H:%M:%S',time.localtime()))

	timecost(time.time()-start_time)